from multiprocessing import Pool

import cv2
import numpy as np
import os
import pydicom as dicom
import pandas as pd
import shutil
import SimpleITK as sitk
from tqdm import tqdm
# 获取dicom文件的pixel_array
def get_pixel_array(dicom_file_path):
    dcm_ori = dicom.dcmread(dicom_file_path)
    arr = dcm_ori.pixel_array
    arr = arr.astype(np.int32)
    return arr


# 最大最小归一化
def pixel_array_normalize(pixel_array):
    return 255.0 * (pixel_array - pixel_array.min()) / (pixel_array.max() - pixel_array.min())


def get_subtraction(array1, array2):
    """
    获取剪影图像
    :param dicom1:  增强后的dicom文件路径
    :param dicom2:  增强前的dicom文件路径
    :return:        剪影图像矩阵
    """
    pa1 = array1
    pa2 = array2
    pa1 = pixel_array_normalize(pa1)
    pa2 = pixel_array_normalize(pa2)
    img = pa1 - pa2

    # 截断负值
    img = np.clip(img, a_min=0, a_max=img.max())
    img = pixel_array_normalize(img)
    # 计算直方图
    img = np.uint8(img)
    hist = cv2.calcHist([img], [0], None, [255], [1, 256])
    # plt.plot(hist)
    # plt.show()
    hist_sum = np.cumsum(hist)
    # 截断归一化
    min_val = 0
    # min_val = np.searchsorted(hist_sum, hist_sum[-1] * 0.01, side='left')
    max_val = np.searchsorted(hist_sum, hist_sum[-1] * 0.9998, side='left')
    img = np.clip(img, a_min=min_val, a_max=max_val)
    img = 255.0 * (img - min_val) / (max_val - min_val)

    return img
def get_center_slice(annotation_path):
    """
    获取每个病例病灶中心切片号
    :param annotation_path: 标注文件路径
    :return: 字典 key为病例ID value为切片号
    """
    annotations = pd.read_excel(annotation_path)
    caseId2sliceId = {}
    for annotation in annotations.itertuples():
        case_id = annotation[1]
        start_slice = annotation[6]
        end_slice = annotation[7]
        caseId2sliceId[case_id] = int((end_slice - start_slice) / 2 + start_slice)
    return caseId2sliceId



# 获取duke数据集的剪影jpg图像
def duke2subtraction(file):
    jpg_save_path = "/data1/home/liukai/AllData/Duke/Duke/myjpg"
    nii_path = "/data1/home/liukai/AllData/Duke/Duke/nii"
    global caseId2sliceId
    pre_dir = os.path.join(nii_path, "pre")
    post2_dir = os.path.join(nii_path, "post2")

    sliceId = caseId2sliceId[file.split('.')[0]]
    pre_nii = sitk.ReadImage(os.path.join(pre_dir, file))
    pre_array = sitk.GetArrayFromImage(pre_nii).astype(np.int32)

    post2_nii = sitk.ReadImage(os.path.join(post2_dir, file))
    post2_array = sitk.GetArrayFromImage(post2_nii).astype(np.int32)
    select_ids = [ sliceId]
    for id in select_ids:
        # 因为数组下标是从0开始，sliceId下标是从1开始，所以对应的数组位置要减1
        pre_center_array = pre_array[id - 1, ::]
        post2_center_array = post2_array[id - 1, ::]
        img_array = get_subtraction(post2_center_array, pre_center_array)
        cv2.imwrite(
            os.path.join(jpg_save_path, file.split('.')[0] + '_second_subtraction_norm_' + str(id) + '.jpg'), img_array)

    print(file)


def draw_box(jpg_dir, img_with_box_save_dir,annotation_path):
    """
    画出标注框
    :param img_dir: 切片图片文件夹路径
    :param annotation_path: 标注文件路径
    :return: None
    """
    if not os.path.exists(img_with_box_save_dir):
        os.makedirs(img_with_box_save_dir)
    annotations = pd.read_excel(annotation_path)
    count=0
    for annotation in annotations.itertuples():
        case_id = annotation[1]
        start_row = annotation[2]
        start_column = annotation[4]
        end_row = annotation[3]
        end_column = annotation[5]
        start_slice = annotation[6]
        end_slice = annotation[7]
        sliceId = int((end_slice - start_slice) / 2 + start_slice)
        if sliceId < 100:
            sliceId = str(sliceId)
        else:
            sliceId = str(sliceId)
        read_name = os.path.join(jpg_dir, case_id + '_second_subtraction_norm_' + sliceId + '.jpg')
        if os.path.exists(read_name):
            print(read_name)
            img = cv2.imread(read_name)
            # img = cv2.cvtColor(img, cv2.COLOR_GRAY2BGR)

            cv2.rectangle(img, (start_column, start_row), (end_column, end_row), (0, 255, 0), 1)
            save_name=os.path.join(img_with_box_save_dir, os.path.basename(read_name))
            cv2.imwrite(save_name, img)
            count+=1
    print(count)


def read_jpg_test():
    jpg_path="/data1/home/liukai/AllData/Duke/Duke/myjpg/Breast_MRI_001_second_subtraction_norm_100.jpg"
    jpg=cv2.imread(jpg_path,flags = 0)#0是灰度，1是白色
    jpg_array=np.array(jpg)
    file="Breast_MRI_001.nii.gz"
    nii_path="/data1/home/liukai/AllData/Duke/Duke/nii"
    pre_dir = os.path.join(nii_path, "pre")
    post2_dir = os.path.join(nii_path, "post2")
    caseId2sliceId = get_center_slice('/data1/home/liukai/AllData/Duke/Duke/Annotation_Boxes.xlsx')
    sliceId = caseId2sliceId[file.split('.')[0]]
    pre_nii = sitk.ReadImage(os.path.join(pre_dir, file))
    pre_array = sitk.GetArrayFromImage(pre_nii).astype(np.int32)

    post2_nii = sitk.ReadImage(os.path.join(post2_dir, file))
    post2_array = sitk.GetArrayFromImage(post2_nii).astype(np.int32)
    select_ids = [sliceId]
    for id in select_ids:
        # 因为数组下标是从0开始，sliceId下标是从1开始，所以对应的数组位置要减1
        pre_center_array = pre_array[id - 1, ::]
        post2_center_array = post2_array[id - 1, ::]
        img_array = get_subtraction(post2_center_array, pre_center_array)
        test=(img_array==jpg_array).all()
        print(test)
if __name__ == '__main__':
    # annotation_path = 'E:\\MRI\\Duke\\Annotation_Boxes.xlsx'
    # get_center_slice(annotation_path)

    # 获取duke数据集的subtraction图像
    # dir_path = 'E:\\MRI\\Duke\\manifest-1680071275430\\Duke-Breast-Cancer-MRI'
    # save_path = 'E:\\MRI\\Duke\\subtractions'
    # duke2subtraction(dir_path, save_path)

    # 切片根据是否含病灶进行分类
    # annotation_path = 'E:\\MRI\\Duke\\Annotation_Boxes.xlsx'
    # save_path = 'E:\\MRI\\Duke'
    # dir_path = 'E:\\MRI\\Duke\\case'
    # duke_classify(annotation_path, dir_path, save_path)

    # ------将duke数据集的中心切片转化为jpg图片并保存---------#
    # dir_path = '/data1/home/liukai/projects/TransFuseForBreast/dataprepare/Duke/dukeExample/Duke-Breast-Cancer-MRI'
    # jpg_save_path = '/data1/home/liukai/projects/TransFuseForBreast/dataprepare/Duke/dukeExample/second_post_center_norm'

    nii_path="/data1/home/liukai/AllData/Duke/Duke/nii"
    caseId2sliceId = get_center_slice('/data1/home/liukai/AllData/Duke/Duke/Annotation_Boxes.xlsx')

    import time
    file_name_list = os.listdir(os.path.join(nii_path, 'pre'))
    file_name_list=["Breast_MRI_015.nii.gz",
                  "Breast_MRI_006.nii.gz","Breast_MRI_001.nii.gz","Breast_MRI_003.nii.gz",]

    start = time.time()
    pool_process = Pool(1)
    pool_process.map(duke2subtraction, file_name_list)
    end = time.time()
    #
    # print('end!cost time:{}/min'.format((end - start) / 60))
    #
    # # -------将jpg图片画上box标注并保存-----#
    # jpg_path = "/data1/home/liukai/AllData/Duke/Duke/myjpg"
    # img_with_box_save_dir='/data1/home/liukai/AllData/Duke/Duke/myjpgWithBox'
    # annotation_path = '/data1/home/liukai/AllData/Duke/Duke/Annotation_Boxes.xlsx'
    # draw_box(jpg_dir = jpg_path,img_with_box_save_dir = img_with_box_save_dir, annotation_path = annotation_path)
    # print("end!")
